package org.apache.spark.simon

import java.util.concurrent.{ExecutorService, Executors, ScheduledExecutorService, TimeUnit}

import org.apache.spark.{SecurityManager, SparkConf}
import org.apache.spark.rpc.{RpcCallContext, RpcEndpointRef, RpcEnv, ThreadSafeRpcEndpoint}

import scala.collection.mutable

/**
 * @author changf
 * @date 2024/1/16 9:15 下午
 */
class Master(val rpcEnv: RpcEnv) extends ThreadSafeRpcEndpoint{
  private val idToWorker: mutable.HashMap[String, WorkerInfo] = new mutable.HashMap[String, WorkerInfo]()


  /**
   * Master扫描检测心跳
   */
  override def onStart(): Unit = {
    val service: ScheduledExecutorService = Executors.newScheduledThreadPool(1)
    service.scheduleAtFixedRate(new Runnable {
      override def run(): Unit = {
        val deadWorkers: Iterable[WorkerInfo] = idToWorker.values.filter(w => {
          System.currentTimeMillis() - w.lastHeartBeatTime > 10000
        })

        deadWorkers.foreach(e=>
          //idToWorker.remove(e.workerId)
          idToWorker-=e.workerId
        )
        println(s"当前活跃的worker数量:${idToWorker.size}")
      }
    },0,15,TimeUnit.SECONDS)
  }

  //override val rpcEnv: RpcEnv = _
  override def receive: PartialFunction[Any, Unit] = {
    case "test" =>{
      println("收到了 test消息")
    }
    case RegisterWorker(rpcEndPointRef,workerId:String,workerMemory:Int,workerCores:Int)=>{
      val workerInfo: WorkerInfo = new WorkerInfo(workerId, workerMemory, workerCores)
      println(s"Master接收到worker发送过来的消息 $workerId $workerCores $workerMemory")
      //idToWorker.put(workerId,workerInfo);
      idToWorker(workerId)=workerInfo
      //向worker返回注册成功的消息
      rpcEndPointRef.send(RegisteredWorker)

    }
    case HeartBeat(workerId)=>{
      val workerInfo: WorkerInfo = idToWorker(workerId)
      //最后一次访问时间
      workerInfo.lastHeartBeatTime=System.currentTimeMillis()

    }
  }

  override def receiveAndReply(context: RpcCallContext): PartialFunction[Any, Unit] = {
    case "ask-msg"=>{
      println("Master 接收到worker发送的同步消息")
      //向worker返回消息
      Thread.sleep(3000)
      context.reply("reply-msg")
    }
  }
}
//伴生对象
object Master {
  def main(args: Array[String]): Unit = {
    //1.创建RpcEnv
    val conf: SparkConf = new SparkConf()
    val manager: SecurityManager = new SecurityManager(conf)
    val rpcEnv: RpcEnv = RpcEnv.create("SparkMaster", "localhost", 8888, conf, manager)
    //2.创建RpcEndpoint
    val master: Master = new Master(rpcEnv)
    //3.RpcEnv 注册RpcEndpoint
    val masterEndpointRef: RpcEndpointRef = rpcEnv.setupEndpoint("Master", master)
    //4.通过返回的RpcEndpointRef发送消息
    masterEndpointRef.send("test")
    //程序挂起 等待退出
    rpcEnv.awaitTermination()



  }
}
